package com.yjb.tenet.boot.sse;

import cn.dev33.satoken.annotation.SaIgnore;
import lombok.extern.slf4j.Slf4j;
import org.springframework.web.bind.annotation.CrossOrigin;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RestController;
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;

import java.io.IOException;

/**
 * @author yinjinbiao
 * @create 2024/4/28 11:54
 */
@CrossOrigin(origins = "*")
@Slf4j
@RestController
@RequestMapping("/sse")
public class SseEmitterController {


    /**
     * 连接，获取SseEmitter对象
     * @param clientId
     * @return
     */
    @SaIgnore
    @GetMapping("/emitter")
    public SseEmitter getSseEmitter(String clientId){
        SseEmitter sseEmitter = new SseEmitter(0L);
        sseEmitter.onCompletion(() -> {
            log.info("onCompletion");
            SseEmitterConstants.sseEmitterMap.remove(clientId);
        });
        sseEmitter.onTimeout(() -> {
            log.info("onTimeout");
            SseEmitterConstants.sseEmitterMap.remove(clientId);
        });
        sseEmitter.onError((e) -> {
            log.info("onError");
            SseEmitterConstants.sseEmitterMap.remove(clientId);
        });
        SseEmitterConstants.sseEmitterMap.put(clientId, sseEmitter);
        return sseEmitter;
    }

    /**
     * 向SseEmitter对象发送数据
     *
     * @param clientId
     * @return
     */
    @RequestMapping("/send")
    public String setSseEmitter(String clientId, String msg) {
        try {
            SseEmitter sseEmitter = SseEmitterConstants.sseEmitterMap.get(clientId);
            if (sseEmitter != null) {
                sseEmitter.send(msg);
            }
        } catch (IOException e) {
            log.error("IOException!", e);
            return "error";
        }
        return "Succeed!";
    }

    /**
     * 将SseEmitter对象设置成完成
     *
     * @param clientId
     * @return
     */
    @RequestMapping("/end")
    public String completeSseEmitter(String clientId) {
        SseEmitter sseEmitter =  SseEmitterConstants.sseEmitterMap.get(clientId);
        if (sseEmitter != null) {
            SseEmitterConstants.sseEmitterMap.remove(clientId);
            sseEmitter.complete();
        }
        return "Succeed!";
    }

    @GetMapping("/sendAll")
    public void sendAll(String msg) {
        for (String key : SseEmitterConstants.sseEmitterMap.keySet()) {
            try {
                SseEmitter sseEmitter = SseEmitterConstants.sseEmitterMap.get(key);
                sseEmitter.send(msg);
            } catch (IOException e) {
                log.error("发送消息错误");
            }
        }
    }
}
